#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import io
import re
import abc
import glob
import qemu
import time
import hivex
import shutil
import pickle
import struct
import parted
import zipfile
import pathlib
import pycdlib
import hashlib
import functools
import subprocess
import configparser


class Util:

    driveC = "C:"
    driveD = "D:"
    driveReserve = "SysReserv"

    partitionTypeMbr = "mbr"
    partitionTypeGpt = "gpt"

    fsTypeFat = "vfat"
    fsTypeNtfs = "ntfs"

    diskPartTableMbr = "mbr"
    diskPartTableGpt = "gpt"

    @staticmethod
    def getListByItem(slist, item=None):
        if item is None:
            return slist
        else:
            assert item in slist
            return [item]

    @staticmethod
    def getRelativePathCaseInsensitive(rootDir, relPath):
        out = subprocess.check_output(["find", rootDir, "-iwholename", os.path.join(rootDir, relPath), "-print", "-quit"], text=True)
        out = out.rstrip("\n")
        out = out[len(rootDir) + 1:]       # remove (rootDir + "/"), so that we return relative path
        return out

    @staticmethod
    def boolToInt01(flag):
        return 1 if flag else 0

    @staticmethod
    def boolTo01(flag):
        return "1" if flag else "0"

    @staticmethod
    def listStartswith(theList, subList):
        return len(theList) >= len(subList) and theList[:len(subList)] == subList

    @staticmethod
    def sizeStrToInt(sizeStr):
        if sizeStr.endswith("M"):
            return int(sizeStr[:-1]) * 1024 * 1024
        elif sizeStr.endswith("G"):
            return int(sizeStr[:-1]) * 1024 * 1024 * 1024
        else:
            return int(sizeStr)

    @staticmethod
    def diskSizeStrToInt(sizeStr):
        if sizeStr.endswith("M"):
            return int(sizeStr[:-1]) * 1000 * 1000
        elif sizeStr.endswith("G"):
            return int(sizeStr[:-1]) * 1000 * 1000 * 1000
        else:
            return int(sizeStr)

    @staticmethod
    def generateSha256(buf):
        sha256 = hashlib.sha256()
        sha256.update(buf)
        return sha256.hexdigest()

    @staticmethod
    def forceDelete(path):
        if os.path.islink(path):
            os.remove(path)
        elif os.path.isfile(path):
            os.remove(path)
        elif os.path.isdir(path):
            shutil.rmtree(path)
        elif os.path.lexists(path):
            os.remove(path)             # other type of file, such as device node
        else:
            pass                        # path does not exist, do nothing

    @staticmethod
    def forceMakeEmptyDir(path):
        if not os.path.exists(path):
            os.mkdir(path)
        elif os.path.islink(path) or not os.path.isdir(path):
            Util.forceDelete(path)
            os.mkdir(path)
        else:
            for fn in os.listdir(path):
                Util.forceDelete(os.path.join(path, fn))

    @staticmethod
    def removeDirContentExclude(dirPath, excludeList):
        for fn in os.listdir(dirPath):
            if fn not in excludeList:
                Util.forceDelete(os.path.join(dirPath, fn))

    @staticmethod
    def touchFile(filepath):
        assert not os.path.exists(filepath)
        f = open(filepath, 'w')
        f.close()

    @staticmethod
    def moveSoleDirContentOut(dirPath):
        # we process only 1 directory in dirPath
        tlist = os.listdir(dirPath)
        if len(tlist) != 1:
            return False

        # we process only 1 directory in dirPath
        soleDirFn = tlist[0]
        soleDirFullfn = os.path.join(dirPath, soleDirFn)
        if os.path.islink(soleDirFullfn) or not os.path.isdir(soleDirFullfn):
            return False

        # deal with name duplication
        if soleDirFn in os.listdir(soleDirFullfn):
            newFn = "_"
            while len(newFn) <= 10 and newFn in os.list(soleDirFullfn):
                newFn += "_"
            assert len(newFn) <= 10
            newFullfn = os.path.join(dirPath, newFn)
            os.rename(soleDirFullfn, newFullfn)
            soleDirFn = newFn
            soleDirFullfn = newFullfn

        # do work
        for fn in os.listdir(soleDirFullfn):
            os.rename(os.path.join(soleDirFullfn, fn), os.path.join(dirPath, fn))
        os.rmdir(soleDirFullfn)
        return True

    @staticmethod
    def saveObj(filepath, obj):
        with open(filepath, 'wb') as fh:
            pickle.dump(obj, fh)

    @staticmethod
    def loadObj(filepath, klass):
        with open(filepath, "rb") as fh:
            return pickle.load(fh)

    @staticmethod
    def saveEnum(filepath, obj):
        Util.saveObj(filepath, obj)

    @staticmethod
    def loadEnum(filepath, klass):
        return Util.loadObj(filepath)

    @staticmethod
    def readUntil(f, chList, max=-1, bInclusive=False):
        if isinstance(chList, str):
            buf = ""
        elif isinstance(chList, bytes):
            buf = b''
        elif isinstance(chList, list):
            if isinstance(chList[0], str):
                buf = ""
            elif isinstance(chList[0], bytes):
                buf = b''
            else:
                assert False
        else:
            assert False

        while True:
            c = f.read(1)
            if c in chList:
                if bInclusive:
                    buf += c
                break
            buf += c
            if max >= 0 and len(buf) >= max:
                break

        return buf

    @staticmethod
    def isInstanceList(obj, *instances):
        for inst in instances:
            if isinstance(obj, inst):
                return True
        return False

    @staticmethod
    def listSparseContain(theList, subList):
        # example: theList [1, 2, 3, 4] subList [1, 3] returns True
        idx = -1
        for x in subList:
            idx2 = theList.index(x)
            if idx2 <= idx:
                return False
            idx = idx2
        return True

    @staticmethod
    def getMountDeviceForPath(pathname):
        pathname = os.path.abspath(pathname)
        buf = subprocess.check_output("mount")
        for line in buf.split("\n"):
            m = re.search("^(.*) on (.*) type ", line)
            if m is not None and m.group(2) == pathname:
                return m.group(1)
        return None

    @staticmethod
    def partedGetFreeSpaceRegion(disk):
        region = None
        for r in disk.getFreeSpaceRegions():
            if r.length <= disk.device.optimumAlignment.grainSize:
                # ignore alignment gaps
                continue
            if region is not None:
                raise Exception("only one free region is allowed")
            region = r
        return region

    @staticmethod
    def partedGetPartitionRegion(disk, partSize="*", fromLeftOrFromRight=True):
        assert fromLeftOrFromRight

        freeRegion = Util.partedGetFreeSpaceRegion(disk)

        # make it same as Windows behaviour, 63 sectors per track, first partition starts from track 1
        pStart = max(freeRegion.start, 63)

        pEnd = freeRegion.end
        if partSize != "*":
            m = re.fullmatch("([0-9]+)(MiB|GiB|TiB)", partSize)
            assert m is not None
            sectorNum = parted.sizeToSectors(int(m.group(1)), m.group(2), disk.device.sectorSize)
            if pEnd < pStart + sectorNum - 1:
                raise Exception("not enough space")
            pEnd = pStart + sectorNum - 1

        return parted.Geometry(device=disk.device, start=pStart, end=pEnd)

    @staticmethod
    def partedInitializeDiskFile(diskFile, partitionLayout):
        partitionTableType, partitionInfoList = partitionLayout
        assert partitionTableType in ["msdos", "gpt"]
        assert len(partitionInfoList) >= 1

        def _addPartition(disk, pType, pRegion):
            if pType == "esp":
                assert partitionTableType == "gpt"
                partition = parted.Partition(disk=disk,
                                             type=parted.PARTITION_NORMAL,
                                             fs=parted.FileSystem(type="fat32", geometry=pRegion),
                                             geometry=pRegion)
                partition.setFlag(parted.PARTITION_BOOT)
            elif pType == "msr":
                # FIXME
                assert partitionTableType == "gpt"
                partition = parted.Partition(disk=disk,
                                             type=parted.PARTITION_NORMAL,
                                             fs=parted.FileSystem(type="fat32", geometry=pRegion),
                                             geometry=pRegion)
            elif pType in ["fat16", "fat16-lba", "fat16-boot", "fat16-lba-boot"]:
                partition = parted.Partition(disk=disk,
                                             type=parted.PARTITION_NORMAL,
                                             fs=parted.FileSystem(type="fat16", geometry=pRegion),
                                             geometry=pRegion)
                if "lba" in pType:
                    partition.setFlag(parted.PARTITION_LBA)
                if "boot" in pType:
                    partition.setFlag(parted.PARTITION_BOOT)
            elif pType in ["ntfs", "ntfs-boot"]:
                partition = parted.Partition(disk=disk,
                                             type=parted.PARTITION_NORMAL,
                                             fs=parted.FileSystem(type="ntfs", geometry=pRegion),
                                             geometry=pRegion)
                if "boot" in pType:
                    partition.setFlag(parted.PARTITION_BOOT)
            else:
                assert False
            disk.addPartition(partition=partition,
                              constraint=disk.device.optimalAlignedConstraint)

        # partitionInfoList => preList & postList
        preList = None
        postList = None
        for i in range(0, len(partitionInfoList)):
            pSize, pType = partitionInfoList[i]
            if pSize == "*":
                assert preList is None
                preList = partitionInfoList[:i]
                postList = partitionInfoList[i:]
        if preList is None:
            preList = partitionInfoList
            postList = []

        # create new disk object
        dev = parted.getDevice(diskFile)
        try:
            disk = parted.freshDisk(dev, partitionTableType)
            try:
                # process preList
                for pSize, pType in preList:
                    _addPartition(disk, pType, Util.partedGetPartitionRegion(disk, pSize))

                # process postList
                for pSize, pType in postList:
                    assert pSize == "*"
                    _addPartition(disk, pType, Util.partedGetPartitionRegion(disk))

                # write to disk file
                disk.commit()

                return (dev, disk)
            except BaseException:
                # pyparted sucks: after guess-and-try, I think this is a usable way to CLEANLY do a clean-up
                del disk
                raise
        except BaseException:
            # pyparted sucks: after guess-and-try, I think this is a usable way to CLEANLY do a clean-up
            dev.getPedDevice().destroy()
            del dev
            raise

    @staticmethod
    def partedEnlargeDiskFile(diskFile, partitionLayout, finalSize):
        assert partitionLayout[1][-1][0] == "*"         # the partition with "*" as size must be the last partition
        assert finalSize > os.path.getsize(diskFile)

        # resize the disk file
        with open(diskFile, "rb+") as f:
            f.truncate(finalSize)

        # resize the partition to the maximum size
        # it's strange that libparted python binding sometimes modifies the partition start position when resizing a partition, so we have to rely on parted command
        last_partition_number = len(partitionLayout[1])
        subprocess.check_call(["parted", "--script", diskFile, "resizepart", str(last_partition_number), "100%"])       # use --script to eliminate warning messages

    @staticmethod
    def floppyCopyBootstrapCode(floppy1_filepath, floppy2_filepath):
        with open(floppy1_filepath, "rb") as f:
            with open(floppy2_filepath, "rb+") as f2:
                f.seek(62)
                f2.seek(62)
                f2.write(f.read(512 - 62))

    @staticmethod
    def usbHddCreate(filepath, sizeStr, mountDir, bMountFs=False):
        partitionLayout = (
            "msdos", [
                ("*", "fat16-lba-boot"),    # mkfs.vfat auto enable/disable LBA, we have no swith to control it
                                            # under this lba partition type, windows can do with both non-LBA and LBA filesystem
            ]
        )

        sz = Util.diskSizeStrToInt(sizeStr)
        with open(filepath, 'wb') as f:
            f.truncate(sz)

        dev, disk = Util.partedInitializeDiskFile(filepath, partitionLayout)
        try:
            partiLoopFileList = Util.mountDiskFile(filepath, partitionLayout, mountDir, partedDisk=disk)
            try:
                assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
                subprocess.check_call(["mkfs.vfat", "-F", "16", partiLoopFileList[0]], stdout=subprocess.DEVNULL)    # mkfs.vfat does not have a quiet option
                if bMountFs:
                    mntDir = partiLoopFileList[0][:len(".loop") * -1] + ".mntdir"
                    os.mkdir(mntDir)
                    try:
                        subprocess.check_call(["fusefat", partiLoopFileList[0], mntDir, "-o", "rw+"],
                                              stdout=subprocess.DEVNULL,                       # fusefat has no quiet option
                                              stderr=subprocess.DEVNULL)                       # fusefat has no quiet option
                        return (partiLoopFileList[0], mntDir)
                    except BaseException:
                        os.rmdir(mntDir)
                        raise
                else:
                    return None
            finally:
                if not bMountFs:
                    Util.umountDiskFile(partitionLayout, mountDir)
        finally:
            # pyparted sucks: after guess-and-try, I think this is a usable way to CLEANLY do a clean-up
            del disk
            dev.getPedDevice().destroy()
            del dev

    @staticmethod
    def usbHddUnMount(filepath, mountDir, loopFile, mntDir):
        partitionLayout = (
            "msdos", [
                ("*", "fat16-lba-boot"),    # mkfs.vfat auto enable/disable LBA, we have no swith to control it
                                            # under this lba partition type, windows can do with both non-LBA and LBA filesystem
            ]
        )

        while True:
            try:
                subprocess.check_call(["umount", mntDir])
                break
            except subprocess.CalledProcessError as e:
                if e.returncode == 32:
                    # means device is busy, it sucks that fuseloop umount can return this result instead of waiting until operation is done
                    print("__wait2__")
                    time.sleep(1)
                else:
                    # as a cleanup operation, we must go on even if there's error
                    break
        os.rmdir(mntDir)

        Util.umountDiskFile(partitionLayout, mountDir)

    @staticmethod
    def fatCheckDirty(devPath):
        # from https://github.dev/docbobo/pyfatfs, https://wiki.osdev.org/FAT

        # definitions
        COMMON_BPB_HEADER = "<3s8sHBHBHHBHHHLL"        # BPB header layout in struct formatted string
        FAT12_EBPB_HEADER = "BBBL11s8s"                # FAT12 EBPB header layout
        FAT16_EBPB_HEADER = FAT12_EBPB_HEADER          # FAT16 EBPB header layout
        FAT32_EBPB_HEADER = "LHHLHH12sBBBL11s8s"       # FAT32 EBPB header layout
        FAT_DIRTY_BIT_MASK = 0x01                      # Dirty bit in FAT header
        FAT16_CLEAN_SHUTDOWN_BIT_MASK = 0x8000         # FAT16 bit mask for clean shutdown bit
        FAT32_CLEAN_SHUTDOWN_BIT_MASK = 0x08000000     # FAT32 bit mask for clean shutdown bit

        # variables
        fatBit = None
        cleanShutdownBitMask = None
        firstFat = None
        BS_Reserved1 = None

        with open(devPath, "rb") as f:
            # get information
            bpb = struct.unpack(COMMON_BPB_HEADER, f.read(struct.calcsize(COMMON_BPB_HEADER)))
            BPB_BytsPerSec = bpb[2]
            BPB_SecPerClus = bpb[3]
            BPB_RsvdSecCnt = bpb[4]
            BPB_NumFATs = bpb[5]
            BPB_RootEntCnt = bpb[6]
            BPB_TotSec16 = bpb[7]
            BPB_FATSz16 = bpb[9]
            BPB_TotSec32 = bpb[12]

            # determine fat type
            if BPB_RootEntCnt != 0 or BPB_FATSz16 != 0:
                total_sectors = BPB_TotSec16 if BPB_TotSec16 != 0 else BPB_TotSec32
                data_sectors = total_sectors - (BPB_RsvdSecCnt + (BPB_NumFATs * BPB_FATSz16))
                total_clusters = data_sectors // BPB_SecPerClus
                if total_clusters < 4085:
                    fatBit = 12
                else:
                    fatBit = 16
            else:
                fatBit = 32

            # get more information
            if fatBit == 12:
                assert False                                # FAT12 does not support dirty bit
            elif fatBit == 16:
                cleanShutdownBitMask = FAT16_CLEAN_SHUTDOWN_BIT_MASK
                ebpb = struct.unpack(FAT16_EBPB_HEADER, f.read(struct.calcsize(FAT16_EBPB_HEADER)))
                BS_Reserved1 = ebpb[22]
                f.seek(BPB_RsvdSecCnt * BPB_BytsPerSec)     # seek to first FAT
                buf = f.read(16 // 8)
                firstFat = struct.unpack("<H", buf)[0]
            elif fatBit == 32:
                cleanShutdownBitMask = FAT32_CLEAN_SHUTDOWN_BIT_MASK
                ebpb = struct.unpack(FAT32_EBPB_HEADER, f.read(struct.calcsize(FAT32_EBPB_HEADER)))
                BS_Reserved1 = ebpb[23]
                f.seek(BPB_RsvdSecCnt * BPB_BytsPerSec)     # seek to first FAT
                buf = f.read(32 // 8)
                firstFat = struct.unpack("<I", buf)[0]
                assert False                                # FIXME: we are not sure fat32 logic is correct
            else:
                assert False

        # check
        dos_dirty = (firstFat & cleanShutdownBitMask) != cleanShutdownBitMask
        nt_dirty = (BS_Reserved1 & FAT_DIRTY_BIT_MASK) == FAT_DIRTY_BIT_MASK
        return dos_dirty or nt_dirty

    @staticmethod
    def ntfsCheckDirty(devPath):
        # from https://flatcap.github.io/linux-ntfs/ntfs/files/index.html

        with open(devPath, 'rb') as f:
            # read the Volume Boot Record (VBR)
            vbr = f.read(512)
            assert vbr[3:11] == b'NTFS    '

            # extract the start cluster of the MFT from the VBR (offset 0x30)
            mft_start_cluster = struct.unpack_from('<Q', vbr, 0x30)[0]

            # get the number of sectors per cluster from the BPB (offset 0x0D)
            sectors_per_cluster = vbr[0x0D]

            # calculate the start sector of the MFT
            mft_start_sector = mft_start_cluster * sectors_per_cluster

            # each MFT record is typically 1KB (1024 bytes)
            mft_record_size = 1024

            # $Volume file's MFT record number is typically 3
            volume_mft_entry_number = 3

            # seek to the start of the $MFT record, jump over the first record ($MFT itself), jump over the previous 2 MFT record
            f.seek(mft_start_sector * 512 + mft_record_size + (volume_mft_entry_number - 1) * mft_record_size)
            volume_mft_entry = f.read(mft_record_size)

            # parse the attribute list in the MFT entry to find the $VOLUME_INFORMATION
            attr_offset = struct.unpack_from('<H', volume_mft_entry, 0x14)[0]

            # traverse all attributes to find the $VOLUME_INFORMATION attribute
            filename = None
            dirty_bit = None
            if True:
                current_offset = attr_offset
                while True:
                    type_code = struct.unpack_from('<I', volume_mft_entry, current_offset)[0]
                    if type_code == 0xFFFFFFFF:
                        # end of attribute list
                        break

                    if type_code == 0x30:
                        # $FILE_NAME (Resident type attribute)
                        filename_len_offset = current_offset + 0x18 + 0x40      # jump over Standard Attribute Header (length 0x18) and go to "Filename length in characters" field (offset 0x40)
                        filename_len = struct.unpack_from('<B', volume_mft_entry, filename_len_offset)[0]

                        filename_offset = current_offset + 0x18 + 0x42
                        filename_bytes = struct.unpack_from('<%ds' % (filename_len * 2), volume_mft_entry, filename_offset)[0]
                        filename = filename_bytes.decode("utf-16")

                    if type_code == 0x70:
                        # $VOLUME_INFORMATION (a Resident type attribute)
                        flags_offset = current_offset + 0x18 + 0x0A             # jump over Standard Attribute Header (length 0x18) and go to "Flags" field (offset 0x0A)
                        flags = struct.unpack_from('<H', volume_mft_entry, flags_offset)[0]
                        dirty_bit = flags & 0x0001
                        break

                    length = struct.unpack_from('<I', volume_mft_entry, current_offset + 4)[0]
                    current_offset += length

            assert filename == "$Volume"
            return (dirty_bit != 0)

    @staticmethod
    def isoExtractEltoritoFloppyImageData(filepath):
        iso = pycdlib.PyCdlib()
        iso.open(filepath)
        try:
            if iso.eltorito_boot_catalog is None:
                return None
            assert iso.eltorito_boot_catalog.validation_entry.platform_id == 0                            # 0 means 8086
            assert iso.eltorito_boot_catalog.initial_entry.boot_indicator == 0x88                         # 0x88 means bootable
            with open(filepath, "rb") as f:
                f.seek(iso.eltorito_boot_catalog.initial_entry.load_rba * 2048)   # sector length is 2048
                if iso.eltorito_boot_catalog.initial_entry.boot_media_type == pycdlib.eltorito.EltoritoEntry.MEDIA_12FLOPPY:
                    return f.read(1200 * 1024)  # read 1.2MB data
                elif iso.eltorito_boot_catalog.initial_entry.boot_media_type == pycdlib.eltorito.EltoritoEntry.MEDIA_144FLOPPY:
                    return f.read(1440 * 1024)  # read 1.44MB data
                elif iso.eltorito_boot_catalog.initial_entry.boot_media_type == pycdlib.eltorito.EltoritoEntry.MEDIA_288FLOPPY:
                    return f.read(2880 * 1024)  # read 2.88MB data
                else:
                    assert False
        finally:
            iso.close()

    @staticmethod
    def pycdlibGetLabel(iso):
        return iso.pvds[0].volume_identifier.decode(iso.pvds[0].encoding).rstrip(" ")

    @staticmethod
    def pycdlibHasEltorito(iso):
        return iso.eltorito_boot_catalog is not None and iso.eltorito_boot_catalog.initial_entry.boot_indicator == 0x88

    @classmethod
    def pycdlibReadBinaryFile(cls, iso, **kwargs):
        cls._pycdlibPathCheckAndReturnType(kwargs)

        fp = io.BytesIO()
        iso.get_file_from_iso_fp(fp, **kwargs)
        fp.seek(0)
        return fp.read()

    @classmethod
    def pycdlibDelTree(cls, iso, **kwargs):
        key = cls._pycdlibPathCheckAndReturnType(kwargs)

        # delete recursively, pycdlib doesn't support this function
        for child in iso.list_children(**kwargs):
            if child.is_dir():
                if child is None or child.is_dot() or child.is_dotdot():
                    continue
                cls.pycdlibDelTree(iso, **{key: iso.full_path_from_dirrecord(child)})
            elif child.is_file():
                iso.rm_file(**{key: iso.full_path_from_dirrecord(child)})
            else:
                assert False
        iso.rm_directory(**kwargs)

    @classmethod
    def pycdlibForceDelete(cls, iso, **kwargs):
        try:
            rec = iso.get_record(**kwargs)
        except pycdlib.pycdlibexception.PyCdlibInvalidInput as e:
            # for udf_path, exception is raised if not found, strange
            if str(e) == 'Could not find path':
                return
            raise

        if rec.is_dir():
            cls.pycdlibDelTree(iso, **kwargs)
        else:
            iso.rm_file(**kwargs)

    @classmethod
    def pycdlibAddDirectory(cls, iso, **kwargs):
        key = cls._pycdlibPathCheckAndReturnType(kwargs)
        if key == "joliet_path":
            keyDict = {
                key: kwargs[key],
                "iso_path": cls._pycdlibToIsoPath(kwargs[key]),
            }
        elif key == "iso_path":
            keyDict = kwargs
        else:
            assert False

        iso.add_directory(**keyDict)

        return keyDict

    @classmethod
    def pycdlibAddFile(cls, iso, fileData, **kwargs):
        key = cls._pycdlibPathCheckAndReturnType(kwargs)
        if key == "joliet_path":
            keyDict = {
                key: kwargs[key],
                "iso_path": cls._pycdlibToIsoPath(kwargs[key]),
            }
        elif key == "iso_path":
            keyDict = kwargs
        else:
            assert False

        iso.add_fp(io.BytesIO(fileData), len(fileData), **keyDict)

        return keyDict

    @staticmethod
    def _pycdlibPathCheckAndReturnType(kwargs):
        assert len(kwargs) == 1
        key = list(kwargs.keys())[0]
        assert key in ["iso_path", "rock_ridge_path", "joliet_path", "udf_path"]
        return key

    @staticmethod
    def _pycdlibToIsoPath(path):
        # convert other type of path into ISO9660 path
        # FIXME: implementation is far from good, need to create a new iso class

        isoFilePath = ""
        for i in range(0, len(path)):
            if re.fullmatch(r"[A-Z0-9_\./]", path[i], re.I):
                isoFilePath += path[i].upper()
            else:
                isoFilePath += "_"

        pathList = isoFilePath.split("/")
        pathList2 = []
        for p in pathList:
            tlist = p.split(".")
            assert len(tlist) in [1, 2]
            if len(tlist[0]) > 8:
                tlist[0] = tlist[0][:6] + "~1"
            if len(tlist) == 2:
                assert len(tlist[1]) <= 3
            pathList2.append(".".join(tlist))
        return isoFilePath

    @staticmethod
    def mountDiskFile(diskFile, partitionLayout, mountDir, allowRoot=False, partedDisk=None):
        if partedDisk is None:
            dev = parted.getDevice(diskFile)
            disk = None
            try:
                disk = parted.newDisk(dev)
                assert len(disk.partitions) == len(partitionLayout[1])
            except BaseException:
                # pyparted sucks: after guess-and-try, I think this is a usable way to CLEANLY do a clean-up
                dev.getPedDevice().destroy()
                del dev
                raise
        else:
            disk = partedDisk
            assert len(disk.partitions) == len(partitionLayout[1])

        mountFileList = []
        try:
            for i, p in enumerate(disk.partitions):
                mountFile = os.path.join(mountDir, "partition%d.loop" % (i))
                Util.touchFile(mountFile)
                cmdList = [
                    "fuseloop",
                    "-O", str(p.geometry.start * disk.device.sectorSize),
                    "-S", str(p.geometry.length * disk.device.sectorSize),
                    diskFile,
                    mountFile,
                ]
                if allowRoot:
                    cmdList += [
                        "-o", "allow_root",
                    ]
                subprocess.run(cmdList, stderr=subprocess.DEVNULL)         # fuseloop prints errors about /etc/mtab
                mountFileList.append(mountFile)
            return mountFileList
        except BaseException:
            for fn in reversed(mountFileList):
                while True:
                    try:
                        subprocess.check_call(["umount", mountFile])
                        break
                    except subprocess.CalledProcessError as e:
                        if e.returncode == 32:
                            # means device is busy, it sucks that fuseloop umount can return this result instead of waiting until operation is done
                            print("__wait1__")
                            time.sleep(1)
                        else:
                            # as a cleanup operation, we must go on even if there's error
                            break
                os.unlink(mountFile)
            raise
        finally:
            if partedDisk is None:
                # pyparted sucks: after guess-and-try, I think this is a usable way to CLEANLY do a clean-up
                del disk
                dev.getPedDevice().destroy()
                del dev

    @staticmethod
    def umountDiskFile(partitionLayout, mountDir):
        for i in range(0, len(partitionLayout[1])):
            mountFile = os.path.join(mountDir, "partition%d.loop" % (i))
            while True:
                try:
                    subprocess.check_call(["umount", mountFile])
                    break
                except subprocess.CalledProcessError as e:
                    if e.returncode == 32:
                        # means device is busy, it sucks that fuseloop umount can return this result instead of waiting until operation is done
                        print("__wait2__")
                        time.sleep(1)
                    else:
                        # as a cleanup operation, we must go on even if there's error
                        break
            os.unlink(mountFile)

    @staticmethod
    def imgCropPureColorBorder(img):
        width, height = img.size
        if width == 0 or height == 0:
            return img

        nLeft, nTop, nRight, nBottom = 0, 0, width - 1, height - 1
        c = img.getpixel((0, 0))

        # move nBottom up, keep a minimum 3 pixel border
        while nBottom > nTop and all([img.getpixel((x, nBottom)) == c for x in range(0, width - 1)]):
            nBottom -= 1
        nBottom = min(nBottom + 3, height - 1)

        # move nTop down, keep a minimum 3 pixel border
        while nTop < nBottom and all([img.getpixel((x, nTop)) == c for x in range(0, width - 1)]):
            nTop += 1
        nTop = max(nTop - 3, 0)

        # we only process up and down border, it's enough for us now
        pass

        return img.crop((nLeft, nTop, nRight + 1, nBottom + 1))

    @staticmethod
    def imgCropRectByExcludingPureColor(img, x, y, color, bStrict=True):
        assert 0 <= x < img.size[0] and 0 <= y < img.size[1]
        assert len(color) == 3

        if len(img.getpixel((x, y))) == 4:
            color = (color[0], color[1], color[2], 255)
        elif len(img.getpixel((x, y))) == 3:
            pass
        else:
            assert False

        if img.getpixel((x, y)) == color:
            if bStrict:
                assert False
            else:
                return None

        # find the left border
        left = x - 1
        while left >= 0:
            if img.getpixel((left, y)) == color:
                break
            left -= 1
        left += 1

        # find the top border
        top = y - 1
        while top >= 0:
            if img.getpixel((x, top)) == color:
                break
            top -= 1
        top += 1

        # find the right border
        right = x + 1
        while right < img.size[0]:
            if img.getpixel((right, y)) == color:
                break
            right += 1

        # find the bottom border
        bottom = y + 1
        while bottom < img.size[1]:
            if img.getpixel((x, bottom)) == color:
                break
            bottom += 1

        return img.crop((left, top, right, bottom))

    @staticmethod
    def imgMatchSubImg(img, subImg, position):
        assert len(img.getpixel((0, 0))) == 3                                   # img has no alpha channel (no transparency)
        assert subImg.size[0] > 0 and subImg.size[1] > 0
        assert subImg.size[0] <= img.size[0] and subImg.size[1] <= img.size[1]  # subImg must be equal or smaller than img

        if position == "left-top":
            xStart, xEnd = 0, subImg.size[0]
            yStart, yEnd = 0, subImg.size[1]
        elif position == "left-bottom":
            xStart, xEnd = 0, subImg.size[0]
            yStart, yEnd = img.size[1] - subImg.size[1], img.size[1]
        elif position == "right-top":
            xStart, xEnd = img.size[0] - subImg.size[0], img.size[0]
            yStart, yEnd = 0, subImg.size[1]
        elif position == "right-bottom":
            xStart, xEnd = img.size[0] - subImg.size[0], img.size[0]
            yStart, yEnd = img.size[1] - subImg.size[1], img.size[1]
        elif position == "center":
            assert False
        else:
            assert False

        for sx, x in enumerate(range(xStart, xEnd)):
            for sy, y in enumerate(range(yStart, yEnd)):
                c = subImg.getpixel((sx, sy))
                if len(c) > 3:
                    if c[3] == 0:
                        # pixel in subImg is transparent, don't compare
                        continue
                    else:
                        assert c[3] == 255
                        c = c[:3]
                if img.getpixel((x, y)) != c:
                    return False
        return True

    @staticmethod
    def imgSearchSubImg(img, subImg, region=None, search_from="left-top"):
        assert len(img.getpixel((0, 0))) == 3                                          # img has no alpha channel (no transparency)

        assert subImg.size[0] > 0 and subImg.size[1] > 0

        if region is not None:
            xStart, yStart, xEnd, yEnd = region
        else:
            xStart, yStart, xEnd, yEnd = 0, 0, img.size[0], img.size[1]
        assert subImg.size[0] <= xEnd - xStart and subImg.size[1] <= yEnd - yStart   # subImg must be equal or smaller than region

        xRange = range(xStart, xEnd - subImg.size[0] + 1)
        yRange = range(yStart, yEnd - subImg.size[1] + 1)
        if search_from == "left-top":
            pass
        elif search_from == "left-bottom":
            yRange = reversed(yRange)
        elif search_from == "right-top":
            xRange = reversed(xRange)
        elif search_from == "right-bottom":
            xRange = reversed(xRange)
            yRange = reversed(yRange)
        else:
            assert False

        for sx in xRange:
            for sy in yRange:
                img2 = img.crop((sx, sy, sx + subImg.size[0], sy + subImg.size[1]))
                if Util.imgMatchSubImg(img2, subImg, "left-top"):
                    return (sx, sy)
        return None

    @staticmethod
    def wgetDownload(url, localFile):
        tmpFile = localFile + ".__download__"
        subprocess.check_call(["wget", "-q", "--show-progress", "-O", tmpFile, url])
        os.rename(tmpFile, localFile)

    @staticmethod
    def wgetDownloadAndExtractZipFiles(url, localFileMap):
        result = subprocess.run(["wget", "-q", "--show-progress", "-O", "-", url], stdout=subprocess.PIPE, check=True)
        with zipfile.ZipFile(io.BytesIO(result.stdout), "r") as zf:
            for srcFn, dstFn in localFileMap.items():
                if dstFn is None:
                    dstFn = srcFn
                with zf.open(srcFn) as f:
                    pathlib.Path(dstFn).write_bytes(f.read())


class TempChdir:

    def __init__(self, dirname):
        self.olddir = os.getcwd()
        os.chdir(dirname)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        os.chdir(self.olddir)


class ActionRunner:

    class PersistStorage(abc.ABC):

        @abc.abstractmethod
        def initGetCurrentActionInfo(self):
            pass

        @abc.abstractmethod
        def getHistoryActionNames(self):
            pass

        @abc.abstractmethod
        def isFinished(self):
            pass

        @abc.abstractmethod
        def isInUse(self):
            pass

        @abc.abstractmethod
        def use(self):
            pass

        @abc.abstractmethod
        def saveActionStart(self, actionName):
            pass

        @abc.abstractmethod
        def saveActionEnd(self, error=None):
            pass

        @abc.abstractmethod
        def saveFinished(self):
            pass

        @abc.abstractmethod
        def unUse(self):
            pass

    class CustomAction(abc.ABC):

        @abc.abstractmethod
        def get_after(self):
            pass

        @abc.abstractmethod
        def get_before(self):
            pass

    def Action(after=[], before=[], _custom_action_name=None, _custom_action=None):
        def decorator(func):
            def wrapper(self, *kargs, **kwargs):
                curActionIndex = self._getActionIndex(wrapper._action_func_name)
                if self._finished is not None:
                    if self._finished == "":
                        raise self._errClass("build already finished")
                    else:
                        raise self._errClass("build already failed, %s" % (self._finished))
                if curActionIndex != self._lastActionIndex + 1:
                    lastActionName = self._actionList[self._lastActionIndex]._action_func_name[len("action_"):] if self._lastActionIndex >= 0 else "None"
                    raise self._errClass("action must be executed in order (last: %s, current: %s)" % (lastActionName, wrapper._action_func_name[len("action_"):]))
                self._persistStorage.saveActionStart(wrapper._action_func_name[len("action_"):])
                try:
                    if wrapper._action is None:
                        func(self, *kargs, **kwargs)
                    else:
                        func(self, wrapper._action_func_name[len("action_"):], wrapper._action, *kargs, **kwargs)
                except BaseException as e:
                    self._finished = str(e)
                    self._persistStorage.saveActionEnd(self._finished)
                    raise
                else:
                    self._lastActionIndex = curActionIndex
                    self._persistStorage.saveActionEnd()
            wrapper._action_func_name = (func.__name__ if _custom_action_name is None else "action_" + _custom_action_name)
            wrapper._action = _custom_action
            wrapper._after = (after if _custom_action is None else _custom_action.get_after())
            wrapper._before = (before if _custom_action is None else _custom_action.get_before())
            return wrapper
        return decorator

    def __init__(self, persistStorage, actionList, customActionFunc, errClass):
        self._persistStorage = persistStorage
        self._actionList = actionList
        self._endAction = actionList[-1]
        self._customActionFunc = customActionFunc
        self._errClass = errClass

        # check self._actionList
        self._assertActions()

        # self._lastActionIndex == -1 if no action has been executed
        self._lastActionIndex = -1
        if len(self._persistStorage.getHistoryActionNames()) > 0:
            actionFuncNameList = [x._action_func_name for x in self._actionList]
            historyActionFuncNameList = ["action_" + x for x in self._persistStorage.getHistoryActionNames()]
            self._lastActionIndex = actionFuncNameList.index(historyActionFuncNameList[-1])
            if self._lastActionIndex == -1:
                raise self._errClass("invalid history actions")
            if not Util.listSparseContain(historyActionFuncNameList[:-1], actionFuncNameList[:self._lastActionIndex]):
                raise self._errClass("invalid history actions")

        # not finished:          self._finished is None
        # successfully finished: self._finished == ""
        # abnormally finished:   self._finished == error-message
        _, err = self._persistStorage.initGetCurrentActionInfo()
        if err is not None:
            self._finished = err
        else:
            self._finished = "" if self._persistStorage.isFinished() else None

        self._persistStorage.use()

    def dispose(self):
        self._persistStorage.unUse()

    def finish(self):
        assert self._finished is None
        assert self._lastActionIndex >= self._actionList.index(self._endAction)
        self._finished = ""
        self._persistStorage.saveFinished()

    def add_custom_action(self, action_name, action, insert_after=None, insert_before=None):
        self.add_custom_actions({action_name: action}, insert_after, insert_before)

    def add_custom_actions(self, action_dict, insert_after=None, insert_before=None):
        for action_name, action in action_dict.items():
            assert re.fullmatch("[0-9a-z_]+", action_name) and "action_" + action_name not in dir(self)
            assert isinstance(action, self.CustomAction)
            assert all([re.fullmatch("[0-9a-z_]+", x) for x in action.get_after()])
            assert all([re.fullmatch("[0-9a-z_]+", x) for x in action.get_before()])

        # convert action object or action name to action index
        if insert_before is not None:
            if insert_before.__class__.__name__ == "method":                        # FIXME
                insert_before = self._actionList.index(insert_before)
            else:
                insert_before = self._getActionIndex("action_" + insert_before)
        if insert_after is not None:
            if insert_after.__class__.__name__ == "method":                         # FIXME
                insert_after = self._actionList.index(insert_after)
            else:
                insert_after = self._getActionIndex("action_" + insert_after)

        # convert to use insert_before only
        if insert_before is not None and insert_after is None:
            pass
        elif insert_before is None and insert_after is not None:
            insert_before = insert_after + 1
        elif insert_before is None and insert_after is None:
            insert_before = len(self._actionList)
        else:
            assert False

        assert self._lastActionIndex < insert_before

        # create new actions and add them to self._actionList
        for action_name, action in action_dict.items():
            func = self.Action(_custom_action_name=action_name, _custom_action=action)(self._customActionFunc)
            func = func.__get__(self)
            exec("self.action_%s = func" % (action_name))
            self._actionList.insert(insert_before, eval("self.action_%s" % (action_name)))
            insert_before += 1

        # do check
        self._assertActions()

    def has_action(self, action_name):
        for action in self._actionList:
            if action._action_func_name == "action_" + action_name:
                return True
        return False

    def add_and_run_custom_action(self, action_name, action):
        self.add_and_run_custom_actions({action_name: action})

    def add_and_run_custom_actions(self, action_dict):
        i = self._lastActionIndex
        for action_name, action in action_dict.items():
            if i == -1:
                self.add_custom_action(action_name, action, insert_before=self._actionList[0])
            else:
                self.add_custom_action(action_name, action, insert_after=self._actionList[i])
            i += 1

        exec("self.action_%s()" % (list(action_dict.keys())[0]))

    def get_progress(self):
        if self._finished is None:
            ret = (self._lastActionIndex + 1) * 100 // len(self._actionList)
            return min(ret, 99)
        else:
            return 100

    def _getActionIndex(self, action_func_name):
        for i, action in enumerate(self._actionList):
            if action._action_func_name == action_func_name:
                return i
        assert False

    def _assertActions(self):
        actionFuncNameList = [x._action_func_name for x in self._actionList]
        for i, action in enumerate(self._actionList):
            assert all(["action_" + x not in actionFuncNameList[:i] for x in action._before])
            assert all(["action_" + x not in actionFuncNameList[i+1:] for x in action._after])


class VirtualUser:

    def __init__(self, vm, exception_class):
        self._vm = vm
        self._eclass = exception_class

    def operate(self, op):
        while True:
            try:
                if op.do(self, self._vm.get_screen_content()):
                    return
            except qemu.qmp.error.QMPError:
                if not self._vm.is_running():
                    raise self._eclass("virtual machine aborts abnormally")
                else:
                    raise
            time.sleep(1)

    def sleep(self, seconds):
        assert seconds > 0
        while self._vm.is_running():
            if seconds <= 0:
                return
            time.sleep(1)
            seconds -= 1
        raise self._eclass("virtual machine aborts abnormally")


class FloppyMount:

    def __init__(self, filePath, mountDir, mode="ro"):
        if mode == "ro":
            pass
        elif mode == "rw":
            mode = "rw+"
        else:
            assert False

        self._mountDir = mountDir
        os.makedirs(self._mountDir, exist_ok=True)

        # fusefat has no quiet option
        # FIXME: fusefat's rw is buggy
        subprocess.check_call(["fusefat", filePath, self._mountDir, "-o", mode], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        subprocess.run(["umount", self._mountDir])


class IsoMount:

    def __init__(self, filePath, mountDir):
        self._mountDir = mountDir
        os.makedirs(self._mountDir, exist_ok=True)
        subprocess.check_call(["fuseiso", filePath, self._mountDir])

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        subprocess.run(["umount", self._mountDir])


class WinFile(abc.ABC):

    @abc.abstractmethod
    def get_content(self):
        pass

    def get_binary_content(self):
        # convert to windows format and write
        return self.get_content().replace("\n", "\r\n").encode("iso8859-1")


class WinRegistry:
    """we can't use hivexregedit because it treat REG_SZ and REG_EXPAND_SZ the same, don't know why it implements like this"""

    _valueTypeList = [
        "",
        "REG_SZ",
        "REG_EXPAND_SZ",
        "REG_BINARY",
        "REG_DWORD",
        "REG_DWORD_BIG_ENDIAN",
        "REG_LINK",
        "REG_FULL_RESOURCE_DESCRIPTOR",
        "REG_RESOURCE_REQUIREMENTS_LIST",
        "REG_QWORD"
    ]

    def __init__(self, mountPoint):
        self._mDir = mountPoint

    def get_users(self):
        ret = []
        with TempChdir(self._mDir):
            if os.path.isdir("Users"):
                dn = "Users"
            elif os.path.isdir("Documents and Settings"):
                dn = "Documents and Settings"
            else:
                assert False
            for fullfn in glob.glob(os.path.join(dn, "*", "NTUSER.DAT")):
                ret.append(os.path.basename(os.path.dirname(fullfn)))
        return ret

    @functools.cache
    def get_local_machine_current_control_set_id(self):
        return self.get_value("HKLM\\SYSTEM\\Select", "Current", valueType="REG_DWORD")

    def get_local_machine_current_control_set_key(self):
        return "HKLM\\SYSTEM\\ControlSet%03d" % (self.get_local_machine_current_control_set_id())

    def add(self, key, valueName, valueType, value):
        assert not self.exists(key, valueName, valueType)
        self.add_or_modify(key, valueName, valueType, value)

    def modify(self, key, valueName, valueType, value):
        assert self.exists(key, valueName, valueType)
        self.add_or_modify(key, valueName, valueType, value)

    def add_or_modify(self, key, valueName, valueType, value):
        hivefile, key = self._mapPathToHive(self._processKey(key))

        h = hivex.Hivex(hivefile, write=True)
        try:
            node = self._getNodeEnsure(h, key)

            # check type
            while True:
                try:
                    val = h.node_get_value(node, valueName)        # raises exception when not found
                except BaseException:
                    break
                valtype = h.value_type(val)[0]
                assert valueType == self._valueTypeList[valtype]
                break

            # construct val structure
            val = self._valueToRegVal(valueType, value)
            v = {"key": valueName, "t": val[0], "value": val[1]}
            h.node_set_value(node, v)
            h.commit(hivefile)
        finally:
            del h

    def delete(self, key, valueName=None):
        # if valueName == None, then delete the key
        # if valueName == "*", then delete all the values in key

        hivefile, key = self._mapPathToHive(self._processKey(key))

        h = hivex.Hivex(hivefile, write=True)
        try:
            if valueName is None:
                # delete key
                node = self._getNode(h, key)
                assert node is not None
                assert node != h.root()

                h.node_delete_child(node)
            else:
                # delete value
                node = self._getNode(h, key)
                assert node is not None

                values = []
                for i in h.node_values(node):
                    if valueName == "*" or valueName == h.value_key(i):
                        continue
                    val = h.value_value(i)
                    v = {"key": h.value_key(i), "t": val[0], "value": val[1]}
                    values.append(v)
                h.node_set_values(node, values)

            h.commit(hivefile)
        finally:
            del h

    def exists(self, key, valueName=None, valueType=None):
        if valueType is not None:
            assert valueName is not None

        hivefile, key = self._mapPathToHive(self._processKey(key))

        h = hivex.Hivex(hivefile, write=False)
        try:
            node = self._getNode(h, key)
            if node is None:
                return False

            if valueName is not None:
                try:
                    h.node_get_value(node, valueName)        # raise exception when not found
                except BaseException:
                    return False

            return True
        finally:
            del h

    def get_subkeys(self, key):
        hivefile, key = self._mapPathToHive(self._processKey(key))
        h = hivex.Hivex(hivefile, write=False)
        try:
            return [h.node_name(x) for x in h.node_children(self._getNode(h, key))]
        finally:
            del h

    def get_value(self, key, valueName, valueType=None):
        # Returns None if not exist

        hivefile, key = self._mapPathToHive(self._processKey(key))

        h = hivex.Hivex(hivefile, write=False)
        try:
            node = self._getNode(h, key)
            if node is None:
                return None

            val = None
            try:
                val = h.node_get_value(node, valueName)
            except BaseException:
                return None

            val = h.value_value(val)

            vType, value = self._regValToValue(val)
            if valueType is not None:
                assert valueType == vType

            return value
        finally:
            del h

    def get_value_name_list(self, key):
        hivefile, key = self._mapPathToHive(self._processKey(key))

        h = hivex.Hivex(hivefile, write=False)
        try:
            node = self._getNode(h, key)
            assert node is not None

            ret = []
            for i in h.node_values(node):
                ret.append(h.value_key(i))
            return ret
        finally:
            del h

    def export_to_file(self, filename, key, valueName=None):
        orikey = key
        hivefile, key = self._mapPathToHive(self._processKey(key))

        h = hivex.Hivex(hivefile, write=False)
        try:
            node = self._getNode(h, key)
            assert node is not None

            cfg = configparser.SafeConfigParser()
            c = 0
            for i in h.node_values(node):
                if valueName is not None and h.value_key(i) != valueName:
                    continue
                val = h.value_value(i)
                valueType, value = self._regValToValue(val)

                cfg.add_section("item%d" % (c))
                cfg.set("item%d" % (c), "key", orikey)
                cfg.set("item%d" % (c), "valueName", h.value_key(i))
                cfg.set("item%d" % (c), "valueType", valueType)
                cfg.set("item%d" % (c), "value", value)
                c = c + 1

            cfg.write(open(filename, "w"))                    # fixme: file with delimiter '\n' in windows disk
        finally:
            del h

    def import_from_file(self, filename):
        cfg = configparser.SafeConfigParser()
        cfg.read(filename)
        for secName in cfg.sections():
            key = cfg.get(secName, "key")
            valueName = cfg.get(secName, "valueName")
            valueType = cfg.get(secName, "valueType")
            value = cfg.get(secName, "value")
            self.add_or_modify(key, valueName, valueType, value)

    def _processKey(self, key):
        if key.startswith("HKLM\\"):
            return key
        if key.startswith("HKU\\"):
            return key

        if key.startswith("HKEY_LOCAL_MACHINE\\"):
            return key.replace("HKEY_LOCAL_MACHINE\\", "HKLM\\", 1)
        if key.startswith("HKEY_USERS\\"):
            return key.replace("HKEY_USERS\\", "HKU\\", 1)

        # note:
        # this class edit offline registry, the following keys do not exist at this time:
        #   HKEY_CURRENT_USER
        #   HKEY_CURRENT_CONFIG
        #   HKEY_LOCAL_MACHINE\\SYSTEM\\CurrentControlSet
        #   HKEY_LOCAL_MACHINE\\SYSTEM\\CurrentControlSetXXX\\Hardware Profiles\\Current
        assert False

    def _valueToRegVal(self, valueType, value):
        # returns (valTypeId, val)

        typeId = self._valueTypeList.index(valueType)
        if valueType == "REG_DWORD":
            assert isinstance(value, int)
            val = struct.pack("<I", value)
        elif valueType == "REG_SZ":
            assert isinstance(value, str)
            val = value.encode("utf_16_le")
        elif valueType == "REG_EXPAND_SZ":
            assert isinstance(value, str)
            val = value.encode("utf_16_le")
        else:
            assert False

        return (typeId, val)

    def _regValToValue(self, val):
        # returns (valueType, value)

        valueType = self._valueTypeList[val[0]]
        if valueType == "REG_DWORD":
            value = struct.unpack("<I", val[1])[0]
        elif valueType == "REG_SZ":
            value = val[1].decode("utf_16_le")
            assert value.endswith("\0")
            value = value[:-1]
        elif valueType == "REG_EXPAND_SZ":
            value = val[1][:-1].decode("utf_16_le")
            assert value.endswith("\0")
            value = value[:-1]
        else:
            assert False

        return (valueType, value)

    def _mapPathToHive(self, key):
        # Copied from /usr/bin/virt-win-reg, only support windows NT currently

        hivefile = None
        while True:
            if key.startswith("HKLM\\SAM\\"):
                hivefile = Util.getRelativePathCaseInsensitive(self._mDir, "WINDOWS/system32/config/sam")
                key = key.replace("HKLM\\SAM\\", "", 1)
                break

            if key.startswith("HKLM\\SECURITY\\"):
                hivefile = Util.getRelativePathCaseInsensitive(self._mDir, "WINDOWS/system32/config/security")
                key = key.replace("HKLM\\SECURITY\\", "", 1)
                break

            if key.startswith("HKLM\\SOFTWARE\\"):
                hivefile = Util.getRelativePathCaseInsensitive(self._mDir, "WINDOWS/system32/config/software")
                key = key.replace("HKLM\\SOFTWARE\\", "", 1)
                break

            if key.startswith("HKLM\\SYSTEM\\"):
                hivefile = Util.getRelativePathCaseInsensitive(self._mDir, "WINDOWS/system32/config/system")
                key = key.replace("HKLM\\SYSTEM\\", "", 1)
                break

            if key.startswith("HKLM\\.DEFAULT\\"):
                hivefile = Util.getRelativePathCaseInsensitive(self._mDir, "WINDOWS/system32/config/default")
                key = key.replace("HKLM\\.DEFAULT\\", "", 1)
                break

#        if key.startswith("HKU\\LocalSystem\\")):
#            sid = "S-1-5-18"
#            hivefile = self._lookupPipOfUserSid($sid) + "/NTUSER.DAT";
#        if key.startswith("HKU\\LocalService\\")):
#            sid = "S-1-5-19"
#            hivefile = self._lookupPipOfUserSid($sid) + "/NTUSER.DAT";
#        if key.startswith("HKU\\NetworkService\\")):
#            sid = "S-1-5-20"
#            hivefile = self._lookupPipOfUserSid($sid) + "/NTUSER.DAT";
#
#        m = re.search("^HKU\\(S-1-5-[0-9]+)\\", key)
#        if m:
#            sid = m.group(1)
#            hivefile = self._lookupPipOfUserSid(sid) + "/NTUSER.DAT";

            m = re.search(r"^HKU\\(.*?)\\", key)
            if m:
                uname = m.group(1)
                if os.path.isdir(os.path.join(self._mDir, "Users", uname)):
                    hivefile = os.path.join("Users", uname, "NTUSER.DAT")
                elif os.path.isdir(os.path.join(self._mDir, "Documents and Settings", uname)):
                    hivefile = os.path.join("Documents and Settings", uname, "NTUSER.DAT")
                else:
                    assert False
                key = key.replace(m.group(0), "", 1)
                break

            assert False

        return (os.path.join(self._mDir, hivefile), key)

    def _getNodeEnsure(self, h, nodePath):
        # Get or create node, always return a valid node

        node = h.root()
        for part in nodePath.split("\\"):
            node2 = h.node_get_child(node, part)
            if node2 is not None:
                node = node2
            else:
                node = h.node_add_child(node, part)
        return node

    def _getNode(self, h, nodePath):
        # Get node, returns None if not found

        node = h.root()
        for part in nodePath.split("\\"):
            node = h.node_get_child(node, part)
            if node is None:
                return None
        return node
